
package com.lr.excel.listener;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.context.AnalysisContext;
import com.alibaba.excel.event.AnalysisEventListener;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.lr.excel.entity.ExcelData;
import org.apache.commons.collections4.CollectionUtils;

import java.util.ArrayList;
import java.util.List;

/**
 * @author liurui
 * @ad Java一对一辅导，手把手带你拿高薪！
 * @WeChat：liuliang70
 * @date 2024/6/14 19:27
 */
public class SplitExcelIntoMultipleSheetListener extends AnalysisEventListener<ExcelData> {
    String excelFilePath = "million_rows_sheet_data.xlsx";

    private static final int SHEET_COUNT = 50000;
    private static int count = 0;
    private final List<ExcelData> cachedDataList = new ArrayList<>();
    private static final int BATCH_COUNT = 1000;
    private static final String SHEET_NAME = "data_";
    private static int sheetIndex = 1;
    ExcelWriter excelWriter =null;
    WriteSheet writeSheet = null;

    @Override
    public void invoke(ExcelData excelData, AnalysisContext analysisContext) {
        initWriteObject();
        count++;
        cachedDataList.add(excelData);
        if (cachedDataList.size() == BATCH_COUNT) {
            save();
            cachedDataList.clear();
        }
        if (count % SHEET_COUNT == 0) {
            save();
            cachedDataList.clear();
            sheetIndex++;
            writeSheet = EasyExcel.writerSheet(SHEET_NAME + sheetIndex).build();
        }
    }

    private void initWriteObject() {
        if (excelWriter == null) {
            excelWriter = EasyExcel.write(excelFilePath, ExcelData.class).build();
        }
        if (writeSheet == null) {
            writeSheet = EasyExcel.writerSheet(SHEET_NAME + sheetIndex).build();
        }
    }

    @Override
    public void doAfterAllAnalysed(AnalysisContext analysisContext) {
        if (!CollectionUtils.isEmpty(cachedDataList)) {
            save();
        }
        // Ensure to close the writer after all data is written
        excelWriter.finish();
    }

    private void save() {
        excelWriter.write(cachedDataList, writeSheet);
    }
}
